import os
import os.path as osp
import numpy as np
import cv2


def get_img_list(path):
    img_list = []
    for root, dirs, files in os.walk(path):
        for f in files:
            if f == 'preds_mask.png':
                img_list.append(os.path.join(root, f))

    assert len(img_list) > 0, "没有图片！"
    return img_list

if __name__ == '__main__':
    # img_list = get_img_list(r"evaluation_logs\v2_res")
    img_list = [r"evaluation_logs\v3_res_post2\Image2022331613913\preds_mask.png"]
    save_dir = r"evaluation_logs/tmp"
    os.makedirs(save_dir, exist_ok=True)
    

    for img_path in img_list:
    # img_path = r'evaluation_logs\tmp\Image202231169943\preds_mask.png'
        img = cv2.imread(img_path)

        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # cv2.imshow('gray', gray)
        # ret, thresh = cv2.threshold(gray, 0, 255,cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        # cv2.imshow('thresh', thresh)
        cv2.imwrite(osp.join(save_dir, 'gray.png'), gray)

        # noise removal
        kernel = np.ones((3,3),np.uint8)
        opening = cv2.morphologyEx(gray, cv2.MORPH_OPEN,kernel, iterations = 5)
        # cv2.imshow('opening', opening)
        cv2.imwrite(osp.join(save_dir, 'opening.png'), opening)

        # sure background area
        sure_bg = cv2.dilate(opening, kernel, iterations=1)
        # cv2.imshow('sure_bg', sure_bg)
        cv2.imwrite(osp.join(save_dir, 'sure_bg.png'), sure_bg)

        # Finding sure foreground area
        dist_transform = cv2.distanceTransform(opening, cv2.DIST_L2, 3)
        # dist_transform_show = (dist_transform / dist_transform.max() * 255).astype(np.uint8)
        # cv2.imshow('dist_transform_show', dist_transform_show)

        ret, sure_fg = cv2.threshold(dist_transform, 0.2 * dist_transform.max(), 255, 0)
        # cv2.imshow('sure_fg', sure_fg)
        cv2.imwrite(osp.join(save_dir, 'sure_fg.png'), sure_fg)

        # Finding unknown region
        sure_fg = np.uint8(sure_fg)
        unknown = cv2.subtract(sure_bg,sure_fg)
        # cv2.imshow('unknown', unknown)
        cv2.imwrite(osp.join(save_dir, 'unknown.png'), unknown)

        # Marker labelling
        ret, markers = cv2.connectedComponents(sure_fg)
        
        # Add one to all labels so that sure background is not 0, but 1
        markers = markers+1

        # Now, mark the region of unknown with zero
        markers[unknown==255] = 0

        markers = cv2.watershed(img, markers)

        bg_mask = np.zeros_like(img)
        bg_mask[markers == -1] = [0,0,255]
        # cv2.imshow('bg_mask', bg_mask)
        cv2.imwrite(osp.join(save_dir, 'bg_mask.png'), bg_mask)

        pix_unique = np.unique(markers)
        area_list = []
        for pix in pix_unique:
            if pix == -1:
                continue
            area = (markers == pix).sum()
            if area > 1600*1600*0.25:
                continue
            area_list.append(area)

        ins_nums = 0
        area_mean = np.mean(area_list)
        for area in area_list:
            if area < area_mean * 0.25:
                continue
            ins_nums += 1
            scale = area / area_mean
            if 1.6 < scale < 2.6:
                ins_nums += 1
            elif scale >= 2.4:
                ins_nums += 2

        for pix in np.unique(markers):
            if pix == -1:
                continue
            b = np.random.randint(0, 255)
            g = np.random.randint(0, 255)
            r = np.random.randint(0, 255)
            img[markers == pix] = [b, g, r]
        
        # cv2.imshow('res', img)
        img_name = img_path.split('\\')[-2]
        # cv2.imwrite('evaluation_logs/{}.png'.format(img_name), img)
        cv2.imwrite(osp.join(save_dir, 'res.png'), img)



        print("{} nums: {}".format(img_name, ins_nums))

        cv2.waitKey(0)